#include<iostream>
#include<set>

using namespace std;

template<class T>
struct rb_tree_node
{
    T _val;
    bool _red;   // 为true说明是红色，否则为黑色
    rb_tree_node* _left;
    rb_tree_node* _right;
    rb_tree_node* _parent;
    rb_tree_node(const T& x = T())
        :_val(x)
        ,_red(true)
        ,_left(nullptr)
        ,_right(nullptr)
        ,_parent(nullptr)
    {}
};

template<class T>
class rb_tree
{
    typedef rb_tree_node<T> node;
    typedef rb_tree_node<T>* node_ptr;
protected:
    node* _root = nullptr;

public:

    node* get_node(const T& val) { return new node(val); }
    void del_node(node* cur) { delete cur; }

    // 找到返回权值为x的节点，否则返回x应该插入的位置的父节点
    node* find(const T& val)
    {
        node* p = nullptr;
        node* cur = _root;
        while(cur)
        {
            p = cur;
            if(val < cur->_val) cur = cur->_left;
            else if(val > cur->_val) cur = cur->_right;
            else return cur;
        }
        return p; 
    }

    bool insert(const T& val)
    {
        if(!_root) 
        {
            _root = get_node(val);
            _root->_red = false;
            return true;
        }
        node* cur = find(val);
        if(cur->_val == val)
            return false;
        else{
            node* p = cur;
            cur = get_node(val);
            cur->_parent = p;
            if(val < p->_val) p->_left = cur;
            else p->_right = cur;
            // 平衡调整
            insert_update(cur);
            return true;
        }
    }
    void insert_update(node* n)
    {
        if(n == _root)
        {
            n->_red = false;
            return;
        }
        node* p = n->_parent;
        if(!p->_red) return;

        node* g = p->_parent;
        node* u = g->_left == p ? g->_right : g->_left;
        // 1.叔叔为红
        if(u && u->_red)
        {
            p->_red = u->_red = false;
            g->_red = true;
            insert_update(g);
        }
        // 2.叔叔为黑（存在且为黑 / 空节点）
        else{
            // 1). 折线型。需转化成直线型
            if((p == g->_left) ^ (n == p->_left))
            {
                if (p == g->_left && n == p->_right)
                    rotateL(p);
                else
                    rotateR(p);
                std::swap(n, p);
            }
            // 2). 直线型
            if(p == g->_left)
                rotateR(g);
            else 
                rotateL(g);
            std::swap(p->_red, g->_red);
            p->_red = false;
        }
    }
    bool check()
    {
        if(!_root) return true;
        if(_root->_red) return false;

        int cnt = 0;
        node* cur = _root;
        while(cur)
        {
            if(!cur->_red) cnt++;
            cur = cur->_left;
        }
        return check_color(_root, 0, cnt);
    }
    bool check_color(node* x, int k, int cnt)
    {
        if(!x) return k == cnt;
        if(x->_red)
        {
            if(x->_left && x->_left->_red) return false;
            if(x->_right && x->_right->_red) return false;
        }
        else k++;
        if(!check_color(x->_left, k, cnt)) return false;
        if(!check_color(x->_right, k, cnt)) return false;
        return true;
    }
    void print()
    {
        _print(_root);
        cout << endl;
    }
    void _print(node* x)
    {
        if(!x) return;
        _print(x->_left);
        cout << x->_val << ' ';
        _print(x->_right);
    }
    //     |               |
    //     p               x
    //    / \             / \
    //   x   C   ====>   A   p
    //  / \                 / \
    // A   B               B   C
    void rotateR(node* p)
    {
        node* g = p->_parent;
        node* x = p->_left;
        p->_left = x->_right;
        if(x->_right) x->_right->_parent = p;

        x->_right = p;
        p->_parent = x;

        x->_parent = g;
        if(g){
            if(p == g->_left) g->_left = x;
            else g->_right = x;
        }
        if(p == _root) _root = x;
    }

    //     |               |
    //     p               x
    //    / \             / \
    //   x   C   <====   A   p
    //  / \                 / \
    // A   B               B   C
    void rotateL(node* x)
    {
        node* g = x->_parent;
        node* p = x->_right;
        x->_right = p->_left;
        if(p->_left) p->_left->_parent = x;

        p->_left = x;
        x->_parent = p;

        p->_parent = g;
        if(g){
            if(x == g->_left) g->_left = p;
            else g->_right = p;
        }
        if(x == _root) _root = p;
    }
};
